from enum import Enum
import io
import json
import os
from pathlib import Path
import re
import tempfile
from matplotlib import pyplot as plt
from flexitext import flexitext
import numpy as np
from pydantic import BaseModel, Field


class UnionBase(Enum):
    # 一个碱基
    A = 'A'
    C = 'C'
    G = 'G'
    T = 'T'
    # 两个碱基
    CT = 'Y'
    AG = 'R'
    AT = 'W'
    CG = 'S'
    GT = 'K'
    AC = 'M'
    # 三个碱基 CGT
    CK = 'B'
    GY = 'B'
    ST = 'B'
    CGT = 'B'
    # 三个碱基 AGT
    AK = 'D'
    RT = 'D'
    GW = 'D'
    AGT = 'D'
    # 三个碱基 ACT
    AY = 'H'
    CW = 'H'
    MT = 'H'
    ACT = 'H'
    # 三个碱基 ACG
    AS = 'V'
    CR = 'V'
    GM = 'V'
    ACG = 'V'
    # 四个碱基 ACG
    ACGT = 'N'
    N = 'N'


class Color(Enum):
    A = 'green'
    T = 'red'
    G = 'black'
    C = 'blue'
    N = 'gray'


class SNV(BaseModel):
    chrom: str
    pos: int
    ref: str
    alt: str
    gene: str
    sample: str

    @classmethod
    def parse_str(cls, text: str) -> 'SNV':
        sample, gene, chrom, pos, ref, alt = re.split(r':|-', text)
        return cls(sample=sample, chrom=chrom, pos=int(pos), ref=ref, alt=alt, gene=gene)

    @property
    def title(self):
        return f'{self.sample} <style:italic>{self.gene}</> {self.chrom}-{self.pos}-{self.ref}-{self.alt}'


class Alignment(BaseModel):
    index: list[int] = Field(alias='pos')
    peak_A: list[int] = Field(alias='peakA')
    peak_C: list[int] = Field(alias='peakC')
    peak_G: list[int] = Field(alias='peakG')
    peak_T: list[int] = Field(alias='peakT')
    basecall_index: list[int] = Field(alias='basecallPos')
    basecall_nts: dict[str, str] = Field(alias='basecalls')
    chrom: str = Field(alias='refchr')
    pos: int = Field(alias='refpos')
    ref_seq: str = Field(alias='refalign')
    alt_seq: str = Field(alias='altalign')
    forward: int

    class Config:
        populate_by_name = True

    @classmethod
    def savgol_filter(cls, data, window_length, polyorder, deriv=0, delta=1.0):
        if window_length % 2 != 1 or window_length <= 0:
            raise ValueError("window_length must be a positive odd integer")
        if polyorder < 0:
            raise ValueError("polyorder must be non-negative")
        if polyorder >= window_length:
            raise ValueError("polyorder must be less than window_length")

        half_window = window_length // 2

        # 计算Vandermonde矩阵
        vander = np.vander(np.arange(-half_window, half_window + 1), polyorder + 1, increasing=True)

        # 计算权重矩阵
        weights = np.linalg.pinv(vander)[deriv] / delta**deriv

        # 对数据进行补零处理
        padded_data = np.pad(data, (half_window, half_window), mode='edge')

        # 应用滤波器
        smoothed_data = np.convolve(padded_data, weights, mode='valid')

        return smoothed_data


    @property
    def smooth_peaks(self):
        peaks = {}
        for nt in ['A', 'T', 'G', 'C']:
            peaks[nt] = self.savgol_filter(getattr(self, f'peak_{nt}'), window_length=11, polyorder=3)
        return peaks

    @property
    def basecall_union_nts(self) -> list[str]:
        nts = list()
        for i in self.basecall_index:
            nt = self.basecall_nts[str(i)]
            if nt != '-':
                nt = nt.split(':')[1]
                nt = getattr(UnionBase,  ''.join(sorted(nt.split('|')))).value
            nts.append(nt)
        return nts

    def calc_variant_index(self, pos: int) -> tuple[int, int, int, int]:
        diff_pos = pos - self.pos
        if self.forward == 0:
            skip = self.alt_seq[::-1][:diff_pos+1].count('-')
            diff_pos = len(self.basecall_index) - diff_pos - 1 - skip
        else:
            skip = self.alt_seq[:diff_pos+1].count('-')
            diff_pos = diff_pos + skip
        index = self.basecall_index[diff_pos-25:diff_pos+25]
        vindex = self.basecall_index[diff_pos]
        return index[0], index[-1], vindex, max([self.peak_A[vindex], self.peak_C[vindex], self.peak_G[vindex], self.peak_T[vindex]])

    def plot_peaks(self):
        for nt, peaks in self.smooth_peaks.items():
            plt.plot(self.index, peaks, color=getattr(Color, nt).value, label=nt)
        plt.legend(loc='upper left')

    def plot_xticks(self, pos: int):
        xticks = plt.xticks(ticks=self.basecall_index, labels=self.basecall_union_nts)
        for i, xtick in enumerate(xticks[1]):
            xtick.set_color(getattr(Color, self.basecall_union_nts[i], Color.N).value)
        xmin, xmax, xvar, yvar = self.calc_variant_index(pos)
        ymax = max(self.peak_A[xmin:xmax] + self.peak_C[xmin:xmax] + self.peak_G[xmin:xmax] + self.peak_T[xmin:xmax]) + 30
        plt.xlim(xmin, xmax)
        plt.ylim(0, ymax)
        plt.arrow(xvar, yvar + 250, 0, -150, head_width=3, head_length=20, lw=5, color='darkred', length_includes_head=True)

    def plot(self, snv: SNV):
        # plt.title(snv.title, loc="center")
        # plt.text(0.5, 1.05, snv.title.replace(snv.gene, f'<i>{snv.gene}</i>'), ha='center', va='center', transform=plt.gca().transAxes)
        flexitext(0.5, 1.05, snv.title, ha='center', va='center')
        plt.ylabel("Peak Density")
        self.plot_peaks()
        self.plot_xticks(snv.pos)


class ABI(BaseModel):
    abi_file: Path
    ref_dir: Path
    snv: SNV
    figsize: tuple = (15, 4)
    dpi:int = 300

    @property
    def reference(self):
        ref = self.ref_dir/f'{self.snv.chrom}.fa.gz'
        if not ref.exists():
            ref = self.ref_dir/f'{self.snv.chrom}.fasta.gz'
        if not ref.exists():
            raise Exception(f'Reference file not found: {ref}')
        return ref

    @property
    def img_name(self):
        return f'{self.snv.sample}.{self.snv.chrom}_{self.snv.pos}.sanger.png'

    def do_align(self) -> Alignment:
        temp_file = tempfile.NamedTemporaryFile(mode='wt')
        temp_file.close()
        cmd = f'tracy align -g {self.reference} {self.abi_file} -o {temp_file.name}'
        if os.system(cmd) > 0:
            raise Exception('Align failed')
        with open(temp_file.name) as reader:
            alignment = Alignment(**json.loads(reader.read()))
            os.unlink(temp_file.name)
            return alignment

    def run(self) -> io.BytesIO:
        alignment = self.do_align()
        bytes_io = io.BytesIO()
        plt.figure(figsize=self.figsize)
        alignment.plot(snv=self.snv)
        plt.savefig(bytes_io, dpi=self.dpi, bbox_inches='tight')
        return bytes_io
